import numpy as np
import platgo as pg
import scipy.io as sio

def Rastrigin(X: np.ndarray) -> np.ndarray:
    X = 0.0512 * X
    return np.sum(X ** 2 - 10 * np.cos(2 * np.pi * X) + 10, axis=1)

def Ackley(X: np.ndarray) -> np.ndarray:
    return -20 * np.exp(-0.2 * np.sqrt(np.mean(X ** 2, axis=1))) - np.exp(np.mean(np.cos(2 * np.pi * X), axis=1)) + 20 + np.exp(1)

def Happycat(X: np.ndarray) -> np.ndarray:
    X = 0.05 * X
    return np.abs(np.sum(X ** 2, axis=1) - X.shape[1]) ** 0.25 + (0.5 * np.sum(X ** 2, axis=1) + np.sum(X, axis=1)) / X.shape[1] + 0.5

def Discus(X: np.ndarray) -> np.ndarray:
    return (1e+6) * X[:, 0] ** 2 + np.sum(X[:, 1:] ** 2, axis=1)

def Rosenbrock(X: np.ndarray) -> np.ndarray:
    X = 0.02048 * X + 1
    return np.sum(100 * (X[:, 0:X.shape[1] - 1] ** 2 - X[:, 1:]) ** 2 + (X[:, 0:X.shape[1] - 1] - 1) ** 2, axis=1)

class CEC_2020_F10(pg.Problem):

    def __init__(self, D=None) -> None:
        self.name = 'CEC_2020_F10'
        self.type['single'], self.type['real'] = [True] * 2
        self.M = 1
        load_path = 'CEC2020.mat'
        load_data = sio.loadmat(load_path)
        mat = []
        for k in load_data.items():
            mat.append(k)
        self.D = D
        self.O = mat[3][1][0][9][0][0][0]
        if self.D is None or self.D < 10:
            self.D = 5
            self.Mat = mat[3][1][0][9][0][0][1]
        elif self.D < 15:
            self.D = 10
            self.Mat = mat[3][1][0][9][0][0][2]
        elif self.D < 20:
            self.D = 15
            self.Mat = mat[3][1][0][9][0][0][3]
        else:
            self.D = 20
            self.Mat = mat[3][1][0][9][0][0][4]
        lb = [-100] * self.D
        ub = [100] * self.D
        self.borders = np.array([lb, ub])
        super().__init__()

    def cal_obj(self, pop: pg.Population) -> None:
        Lambda = np.array([10, 1, 10, (1e-6), 1])
        delta = np.array([10, 20, 30, 40, 50])
        bias = np.array([0, 100, 200, 300, 400])
        W = np.zeros((pop.decs.shape[0], 5))
        F = np.zeros(W.shape)
        for i in range(W.shape[1]):
            tmp = np.sum((pop.decs - np.tile(self.O[i, 0: pop.decs.shape[1]], (pop.decs.shape[0], 1))) ** 2, axis=1)
            W[:, i] = 1 / (np.sqrt(tmp) + (1e-10)) * np.exp(-tmp / 2 / self.D / delta[i] ** 2)
            if i == 0:
                F[:, i] = Rastrigin(np.dot((pop.decs - np.tile(self.O[i, 0:pop.decs.shape[1]], (pop.decs.shape[0], 1))),
                                         self.Mat[0: self.D, :]))
            if i == 1:
                F[:, i] = Happycat(np.dot((pop.decs - np.tile(self.O[i, 0:pop.decs.shape[1]], (pop.decs.shape[0], 1))),
                                        self.Mat[i * self.D:(i + 1) * self.D, :]))
            if i == 2:
                F[:, i] = Ackley(np.dot((pop.decs - np.tile(self.O[i, 0:pop.decs.shape[1]], (pop.decs.shape[0], 1))),
                                          self.Mat[i * self.D:(i + 1) * self.D, :]))
            if i == 3:
                F[:, i] = Discus(np.dot((pop.decs - np.tile(self.O[i, 0:pop.decs.shape[1]], (pop.decs.shape[0], 1))),
                                          self.Mat[i * self.D:(i + 1) * self.D, :]))
            if i == 4:
                F[:, i] = Rosenbrock(np.dot((pop.decs - np.tile(self.O[i, 0:pop.decs.shape[1]], (pop.decs.shape[0], 1))),
                                        self.Mat[i * self.D:(i + 1) * self.D, :]))
        W = W / np.tile(np.sum(W, axis=1).reshape(np.sum(W, axis=1).shape[0], 1), (1, W.shape[1]))
        pop.objv = 2500 + np.sum(W * (np.tile(Lambda, (F.shape[0], 1)) * F + np.tile(bias, (F.shape[0], 1))), axis=1)
        pop.objv = pop.objv.reshape(pop.objv.shape[0], 1)

    def get_optimal(self) -> np.ndarray:
        pass


if __name__ == '__main__':
    problem = CEC_2020_F10()
    alg = pg.algorithms.GA(problem=problem, maxgen=100)
    pop = alg.go(100)
    print(pop)

